"""
Author: Administrator
Date: 2025/1/13
Description:  文件加密工具
"""

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import binascii
import shutil
from Config import AppConfig

def encrypt_data(key: str, model_path: str) -> tuple[str, bytes]:
    """加密模型数据

    使用 AES 算法（高级加密标准）加密给定的模型数据。AES 是对称加密算法，意味着加密和解密使用相同的密钥。
    该函数使用 AES 的 CBC 模式对模型文件进行加密。CBC（Cipher Block Chaining）是分组密码的工作模式之一，
    它通过将每个加密块与前一个块的加密结果进行异或操作，从而增强了加密的安全性。

    :param str key: 密钥1，AES 算法的加密和解密过程需要用到的密钥。
    :param str model_path: 模型文件路径，表示需要加密的文件。
    :return tuple[str, bytes]: 返回密钥2（IV）和加密后的数据。密钥2（IV）是初始化向量，用于解密过程中的安全性。
    """
    # 读取模型文件为二进制数据
    with open(model_path, "rb") as f:
        model_data = f.read()

    # AES.new 创建一个新的 AES 加密对象。这里我们指定了加密模式为 CBC。
    # key.encode() 是将密钥字符串转换为字节，密钥长度应为 16（AES-128）、24（AES-192）或 32（AES-256）字节。
    cipher = AES.new(key.encode(), AES.MODE_CBC)

    # 使用 Crypto.Util.Padding 中的 pad 函数将数据填充到 AES.block_size 的倍数。
    # AES 加密的每个加密块大小是固定的（128 位 = 16 字节），因此需要对原数据进行填充，直到数据大小是 16 字节的倍数。
    padded_data = pad(model_data, AES.block_size)

    # 加密数据。AES 加密操作是块加密，在 CBC 模式下，每个块会通过异或操作与前一个加密块的结果连接在一起。
    # 首先加密第一个块时，由于没有前一个块，它会使用一个随机生成的初始化向量（IV）。
    encrypted_data = cipher.encrypt(padded_data)

    # 获取生成的 IV（初始化向量），它与加密的数据一起传输，用于解密时的对称性保证。
    iv = binascii.hexlify(
        cipher.iv
    ).decode()  # 将 IV 转换为十六进制字符串（便于传输和存储）

    # 返回初始化向量和加密后的数据
    return iv, encrypted_data

def encrypt_model(unencrypted_model_path:str,encrypted_model_path:str):
    # 给定密钥
    key = "ActionEncryptedFile@1029"
    # 加密后的模型文件路径
    if(encrypted_model_path is None):
        encrypted_model_path = r"D:\data\android\demo\aidc\x1ScannerApp\app\src\main\assets\encrypted_model_files\qr/qr_predV9_4d_best_01_quant_perpared_qcs8550.dll"

    # 对模型文件进行加密处理
    iv, encrypted = encrypt_data(key, unencrypted_model_path)

    # 将加密的二进制文件写入dll
    with open(encrypted_model_path, "wb") as f:
        f.write(encrypted)

    print(f"IV: {iv}")
    return iv


# if __name__ == "__main__":
    """
    功能描述： 运行此脚本，对量化后的dlc文件做加密处理，生成加密后的.dll模型文件。
            加密完成后，请将加密后的文件放在X1项目assets资源目录下的encrypted_model_files文件夹下，
            并将密钥2和加密后的文件名写入X1配置文件(ModelFileDecryptionIV.kt 和 ModelFile.kt)中。

    输入： 密钥1、模型文件路径
    输出： 密钥2、加密后的模型文件
    """
    # # 给定密钥
    # key = "ActionEncryptedFile@1029"
    # # 需要被加密的模型文件路径
    # unencrypted_model_path = r"D:\data\android\demo\aidc\x1ScannerApp\app\src\main\assets\encrypted_model_files\qr/qr_predV9_4d_best_01_quant_perpared_qcs8550.dlc"
    # # 加密后的模型文件路径
    # encrypted_model_path = r"D:\data\android\demo\aidc\x1ScannerApp\app\src\main\assets\encrypted_model_files\qr/qr_predV9_4d_best_01_quant_perpared_qcs8550.dll"

    # # 对模型文件进行加密处理
    # iv, encrypted = encrypt_data(key, unencrypted_model_path)

    # # 将加密的二进制文件写入dll
    # with open(encrypted_model_path, "wb") as f:
    #     f.write(encrypted)

    # print(f"IV: {iv}")
    # encrypt_model(AppConfig.model_path,AppConfig.encrypted_model_path)